#!/usr/bin/env perl

use warnings;
use strict;
use 5.012;
use autodie qw(:all);
use Fcntl qw(:seek);

use YAML qw(Dump LoadFile DumpFile); #for debugging only

use File::Temp qw(:POSIX);

use Getopt::Std;
local $Getopt::Std::STANDARD_HELP_VERSION = 1;

$ENV{LC_ALL} = "en_US.UTF8"; # localized versions of mkvtools mess up our regexes
my $EDITOR = ($ENV{EDITOR} ? $ENV{EDITOR} : "vi");

sub HELP_MESSAGE {
    print <<EOF;
Usage:
mkvremux (MP4|M2TS) [OPTIONS] <movie.mkv> [<other files>]

Options:
-s			Split output in files of at most 4GB. Useful for FAT32 filesystems. It tries to
			split in about equal pieces (UNTESTED)
-o <output>	Save output to given path instead of <movie.m2ts>
-m <options> Additional options for MP4Box
-t <options> Additional options for TSMuxeR
-v			Verbose
-e			Edit muxing command line before execution, using \$EDITOR (or vi if unset)
-i			Keep intermediate files (parsing mkvinfo, converted tracks) and try to reuse them
-l			Reduce AVC level from {Main,Extended,High}\@5.1 to {Main,Extended,High}\@4.1,
			as required for PS3 playback
-d			Enable delay checking. When picture and sound are out of sync, it could be that
			the --sync or --delay option (?) of mkvmerge was chosen. Unfortunately, this
			doesn’t get written in the header, so we have to call mkvinfo with the -v switch,
			which basically outputs info on each framegroup. (EXPERIMENTAL)
-p			Try to preserve tracks, e.g., try to include both the converted AC3 track AND 
			the original DTS track. Works currently with m2ts only. In the future, MP4Box
			may	add support for DTS tracks. (EXPERIMENTAL)
EOF
}

sub VERSION_MESSAGE {
    say "=== mkvremux version 0.93 ===";
}
&VERSION_MESSAGE;

my $mp4box = "MP4Box";
my $tsmuxer = "tsMuxeR";
my $mkvextract = "mkvextract";
my $mkvinfo = "mkvinfo";
my $aften = "aften";
my $dcadec = "dcadec";

my %needs_conversion;
$needs_conversion{M2TS} = {"A_AC3" => 0,	# add here conversion subs (see also below)
			"A_DTS" => \&DTS2AC3,			# 0 means: no conv needed
			"V_MPEG4/ISO/AVC" => 0,
			"S_TEXT/UTF8" => 0};

$needs_conversion{MP4} = {"A_AAC" => 0,
			"A_AC3" => 0,
			"A_DTS" => \&DTS2AC3,
			"V_MPEG4/ISO/AVC" => 0,
			"S_TEXT/UTF8" => 0,
			"S_VOBSUB" => 0};
#MP4Box needs correct extensions
my %extension = (	"A_AAC" => "aac",
			"A_AC3" => "ac3",
			"A_DTS" => "dts",		
			"V_MPEG4/ISO/AVC" => "h264",
			"S_TEXT/UTF8" => "srt",
			"S_VOBSUB" => "idx");

my %handler = ("mkv" => \&processMKV,
#NB: one needs to indicate the .idx, not the .sub file!
				"idx" => \&processVobSub); 
				
		
my $mode = shift;
for ($mode) {
	when ("MP4") {
		say "\nOutput mode is MP4.";
	}
	when ("M2TS") {
		say "\nOutput mode is M2TS.";
	}
	default {
		say "\nNo valid output mode selected, quitting…";
		&HELP_MESSAGE && die;
	}
}

my %opts;
getopts('adeilo:ps:v', \%opts)
    or &HELP_MESSAGE && die;


my %movie; # This hash contains all info needed for finding the various tracks and remuxing them
my $tracks = \@{$movie{tracks}}; 	# shortcut so we can say @$tracks instead of @{$movie{tracks}};
				# @$tracks AoH will contain all tracks from the input file and subsequently all tracks which are to be included in the outfile
my $chapters = \@{$movie{chapters}};

my @unlink; # array of temp files to be unlinked after finish

##########################################################################################
################################# Input processing #######################################
##########################################################################################
say "\n= Entering Input Processing =";

foreach my $file (@ARGV) {
	die "\nInput file \"$file\" doesn’t exist!" unless -e $file;

	$file =~ /(.+)\.([[:alnum:]]+)$/ or die "Filename has no extension";
	# The first file determines the basename
	unless (defined $movie{basename}) {
		$movie{basename} = $1;
		if ($opts{i} && -e "$movie{basename}.yaml") {	
			say "Found \"$movie{basename}.yaml\", skipping the parsing of mkvinfo.";
			%movie = %{LoadFile "$movie{basename}.yaml"};
			$tracks = \@{$movie{tracks}}; # redefine $tracks to point to the loaded hash
			$chapters = \@{$movie{chapters}};
			# stop further input processing
			last;
		}
	}
	my $extension = $2;
	$movie{est_filesize} += -s $file; # approx total filesize, for splitting
	say "\nProcessing file \"$file\".";
	$handler{$extension}->($file); #will modify $tracks and $chapters
}

for ($movie{videotrack}{AVC_level}) {
	when ($_ < 5) {
		#FIXME ugly
		say "(First) video track is AVC $movie{videotrack}{AVC_profile} profile, level $movie{videotrack}{AVC_level}.";
		$opts{l} && say "No level switching needed";
		$opts{l} = 0;
	}
	when ($_ >= 5) {
		say "(First) video track is AVC $movie{videotrack}{AVC_profile} profile, level $movie{videotrack}{AVC_level}.";
		$opts{l} && say "Will attempt level switching";
	}
	default {
		say "(First) video track seems to be not AVC encoded" . ($opts{l}? ", turning level switching off" : "");
		$opts{l} = 0;
	}
}

$opts{v} && say "Filesize is $movie{est_filesize} bytes, duration is $movie{duration} seconds.";
$opts{v} && say "Info about input after mkvinfo parsing:\n" . Dump(\%movie);

$opts{i} && DumpFile "$movie{basename}.yaml", \%movie;

## Input handler: mkv
sub processMKV {
	my $inputfile = shift;
	
	my $v =  $opts{d} ? "-v" : ""; # delay checking
	say "Launching $mkvinfo $v $inputfile";
	open(MKVINFO, "-|:encoding(UTF-8)", $mkvinfo, $v, $inputfile);
	
	my $currtrack; # ref to a hash containing info about current track
	my $currchap; # ref to a hash containing info about current chapter

	while(<MKVINFO>) {
		print;
		/ Track number: \d+ \(track ID for mkvmerge & mkvextract: (\d+)\)/ && do {
			# for some (undoubtedly bad) reason the mkv* suite introduced two flavors of track numbers, 0-based and 1-based
			# I decided to use the 0-based ones.
			$currtrack = { num => $1, path => $inputfile };
			push @$tracks, $currtrack;
			next
		};
		/ Track type: (\w+)/ && do {
			$$currtrack{type} = $1;
			# set delay for subtitles to 0, since they contain timing info
			# if they _are_ shifted, we can’t find out with this method
			$$currtrack{delay} = 0 if $$currtrack{type} eq "subtitles"; 
			next
		};
		/ Codec ID: ([\w\/]+)/ && do { $$currtrack{codec} = $1; next };
		/ Language: (\w+)/ && do { $$currtrack{lang} = $1 unless $1 eq "und"; next };
		/ \(([\d.]+) frames\/fields per second for a video track\)/ && ($$currtrack{type} eq "video") && do {
			$$currtrack{fps} = $1;
			$movie{videotrack} = $currtrack unless defined $movie{videotrack}; #we assume here that the first video track is main video track
			next
		};
		/ Pixel width: (\d+)/ && do { $$currtrack{width} = $1; next };
		/ Pixel height: (\d+)/ && do { $$currtrack{height} = $1; next };
		/\(h.264 profile: (\w+) \@L([\d.]+)\)/ && do { 
			$$currtrack{AVC_profile} = $1;
			$$currtrack{AVC_level} = $2;
			next };
		/track number (\d+), [^,]+, timecode ([\d.]+)s/ && do {
			my $track = $1 - 1; # we want 0-based track numbering
			my ($elem) = grep {$$_{num} == $track} @$tracks;
			$$elem{delay} = $2 * 1000 unless defined $$elem{delay}; # want delay in ms
			last if scalar(@$tracks) == grep {defined $$_{delay}} @$tracks;
			next
		};
		/ Duration: (\d+\.\d+)s/ && do { $movie{duration} = $1; next };
		/ ChapterTimeStart: (\d+:\d\d:\d\d\.\d\d\d)/ && do {
			# NB: want only first 3 after comma digits (e.g., VLC gets confused with too
			# digits in chapters)
			$currchap = { start => $1 };
			push @$chapters, $currchap;
			next
		};
		/ ChapterString: (.+)$/ && do {
			$$currchap{name} = $1;
			next
		};

	}
	{no autodie; close MKVINFO;} # we interrupt mkvinfo before we read all of its output
}
## Input handler: VobSub
sub processVobSub {
	my $inputfile = shift;
	
	push @$tracks, {path => $inputfile, 
					num => 0,
					type => "subtitles",
					codec => "S_VOBSUB",
					delay => 0};
}
	

say "\n= Exiting Input Processing =";
##########################################################################################
################################# Track conversion #######################################
##########################################################################################
say "\n= Entering Track Conversion =";
for (@$tracks) {
	unless (defined $needs_conversion{$mode}{$$_{codec}}) {
		print STDERR "\nCannot handle codec \"$$_{codec}\" in $$_{path}:$$_{num}, skipping...";
		#FIXME we don’t really "skip", we just leave it as is, letting it fail later
		# we should remove this track from @$tracks instead
		next;
	}
	if ($needs_conversion{$mode}{$$_{codec}}) {
		$needs_conversion{$mode}{$$_{codec}}->($_); # pass conversion sub a ref to current track
	}
}

$opts{v} && say "Info about input after track conversion:\n" . Dump(\%movie);

$opts{i} && DumpFile "$movie{basename}.yaml", \%movie;

### Conversion subs ###
# Subs get a reference to the track infohash to be converted, which they MAY modify. They MAY also push onto @$track, but MAY NOT splice or unshift it
sub DTS2AC3 {
	my $currtrack = shift;	
	say "Converting DTS track $$currtrack{num} of \"$$currtrack{path}\" to AC3";
	my $dtsfile = tmpnam(); # $dtsfile is a path to a FIFO, through which we pipe the output of mkvextract to dcadec
	my $ac3file = $movie{basename} . "_$$currtrack{num}.$extension{A_AC3}"; # ac3file is the path to the converted file
	if ($opts{i} and -e $ac3file) {
		say "\"$ac3file\" exists already, skipping...";
	} else {
		system "mkfifo", $dtsfile;
		push @unlink, $dtsfile;
		
		my $dcadec_aften_cmd = "$dcadec -o wavall \"$dtsfile\" 2>/dev/null | $aften -v 0 -readtoeof 1 - \"$ac3file\" &";
		my $extract_dts_cmd = "$mkvextract tracks \"$$currtrack{path}\" $$currtrack{num}:\"$dtsfile\"";
		
		$opts{v} && say "Now about to execute \n\t$dcadec_aften_cmd";
		
		system $dcadec_aften_cmd;
		
		$opts{v} && say "Now about to execute \n\t$extract_dts_cmd";
		
		system $extract_dts_cmd;
		
		say "\nDone!";
	}
	if ($opts{p}) {
		push @$tracks, { codec => "A_AC3",
				path => $ac3file,
				lang => $$currtrack{lang},
				delay => $$currtrack{delay},
				num => 0 };
		$movie{est_filesize} += 80000 * $movie{duration}; # AC3 takes with 5.1 channels about 80kbps (???)
	} else {
		$$currtrack{codec} = "A_AC3";
		$$currtrack{path} = $ac3file;
		$$currtrack{num} = 0;
	}

	push @unlink, $ac3file unless $opts{i};
}
say "= Exiting Track Conversion =";
##########################################################################################
################################ Output processing #######################################
##########################################################################################
say "\n= Entering Output Processing =";
my $split_dur;
if ($opts{s}) {
	$movie{est_filesize} += 100 * 1024 ** 2; # add 100MB to estimated filesize for safety
	my $num_pieces = 1 + int $movie{est_filesize} / (4 * 1024 ** 3);
	say "Estimated filesize is $movie{est_filesize}. Seems like we have to split our file in $num_pieces pieces";
	$split_dur = 5 + int $movie{duration} / $num_pieces;
}

### M2TS ###
#tsMuxeR can parse mkv files, so we don’t need to extract all tracks in advance,
#and it handles also the AVC level changing kludge by itself
if ($mode eq "M2TS") {
my @meta; # This array contains all lines of the future .meta file
my %append_meta = (	"A_AC3" => \&meta_audio, # these subs are responsible for writing the .meta file
			"A_DTS" => \&meta_audio,
			"V_MPEG4/ISO/AVC" => \&meta_video,
			"S_TEXT/UTF8" => \&meta_subtitles);

my $maxdelay = 0;
if ($opts{d}) { # find max delay of all audio tracks
	for (grep {$$_{type} eq "audio"} @$tracks) {
		$maxdelay = $$_{delay} if $$_{delay} > $maxdelay;
	}
}

push @meta, "MUXOPT --no-pcr-on-video-pid --new-audio-pes --vbr --vbv-len=500 "
	. ($opts{s} ? "--split-duration=$split_dur " : "")
	. ($maxdelay > 0 ? "--cut-start=${maxdelay}ms " : "")
	. ($opts{t} ? "$opts{t}" : "");

for (@$tracks) {
	$append_meta{$$_{codec}}->($_); # pass .meta generating sub a ref to current track
}

if ($opts{v}) {
	say "Content of .meta file:\n" . join "\n", @meta;
}
my $metafile = $movie{basename} . ".meta";
say "Writing \"$metafile\"...";
open META, ">", "$metafile";
print META join "\n", @meta;
close META;
push @unlink, $metafile;

my $m2tsfile = $movie{basename} . ".m2ts";

$opts{o} =~ s/(.*?)(\.m2ts)?$/$1.m2ts/ if $opts{o}; 

my $tsmuxer_cmd = "$tsmuxer \"$metafile\" " . ($opts{o} ? "\"$opts{o}\"" : "\"$m2tsfile\"");

$opts{e} && editVar(\$tsmuxer_cmd);
$opts{v} && say "Now about to execute \n\t$tsmuxer_cmd";
system $tsmuxer_cmd;

# .meta output subs
# These subs take a ref to a track and push a string in .meta syntax onto @meta
sub meta_audio {
	my $currtrack = shift;
	my $track = $$currtrack{num} + 1; #tsMuxeR wants 1-based numbering
	$opts{v} && say "Appending an audio track with codec \"$$currtrack{codec}\" from file:track \"$$currtrack{path}:$$currtrack{num}\" to .meta file";
	push @meta, "$$currtrack{codec}, \"$$currtrack{path}\", track=$track, lang=$$currtrack{lang}, timeshift=$$currtrack{delay}ms";
}
	
sub meta_video {
	my $currtrack = shift;
	my $track = $$currtrack{num} + 1; #tsMuxeR wants 1-based numbering
	$opts{v} && say "Appending a video track with codec \"$$currtrack{codec}\" from file:track \"$$currtrack{path}:$$currtrack{num}\" to .meta file";
	push @meta, "$$currtrack{codec}, \"$$currtrack{path}\", track=$track, fps=$$currtrack{fps}, insertSEI, contSPS, ar=As source"
		. ($opts{l} ? ", level=4.1" : "");
}

sub meta_subtitles {
	my $currtrack = shift;
	my $track = $$currtrack{num} + 1; #tsMuxeR wants 1-based numbering
	$opts{v} && say "Appending a subtitle track with codec \"$$currtrack{codec}\" from file:track \"$$currtrack{path}:$$currtrack{num}\" to .meta file";
	my $font = $opts{f} ? $opts{f} : "/usr/share/fonts/truetype/Efont_Serif.ttf"; # if someone knows a better way to do this, please help!
	say "Font is $font";
	push @meta, "$$currtrack{codec}, \"$$currtrack{path}\", track=$track, lang=$$currtrack{lang}, font-name=\"$font\", font-size=65, font-color=0x00ffffff, bottom-offset=24, font-border=2, text-align=center, video-width=$movie{videotrack}{width}, video-height=$movie{videotrack}{height}, fps=$movie{videotrack}{fps}";  
}
}
### MP4 ###
if ($mode eq "MP4") {
	my $mp4boxcmd = "$mp4box -new "
		. ($opts{s} ? "-split $split_dur " : "")
		. ($opts{m} ? "$opts{m} " : "");
	my $mkvextractcmd;
	for (@$tracks) {
		if ($$_{path} =~ /.mkv$/) {
			my $outfile = $movie{basename} . "_$$_{num}.$extension{$$_{codec}}";
			
			if ($opts{i} and -e $outfile) {
				say "\"$outfile\" exists already, skipping...";
			} else {
				$mkvextractcmd = "$mkvextract tracks \"$$_{path}\" " unless defined $mkvextractcmd;
				$mkvextractcmd .= "$$_{num}:\"$outfile\" ";
				push @unlink, $outfile unless $opts{i};
			}
			$$_{path} = $outfile;
			$$_{num} = 1;
		}
		$mp4boxcmd .= "-add \"$$_{path}"
			. (exists $$_{fps} ? ":fps=$$_{fps}" : "")
			. ($opts{d} && defined $$_{delay} ? ":delay=$$_{delay}" : "")
			. (exists $$_{lang} ? ":lang=$$_{lang}" : "")
			. "\" ";
	}
	if (defined $mkvextractcmd) {
		#$opts{e} && editVar(\$mkvextractcmd);
		$opts{v} && say "Now about to execute \n\t$mkvextractcmd";
		system $mkvextractcmd;
	}
	## Change AVC level
	if ($opts{l}) {
		print "Trying to change AVC level in $movie{videotrack}{path}…\t";
		open(H264FILE, "+<:raw", $movie{videotrack}{path}); #TODO add support for more than one video track
		seek(H264FILE, 7, SEEK_CUR);
		my $byte;
		read(H264FILE,$byte,1) == 1 or die;
		for(ord($byte)) {
			when([51, 50]) {# 0x33 = 51
				seek(H264FILE, -1, SEEK_CUR);
				print H264FILE "\x{29}";
				say "done.";
			}
			when($_ < 50) { say "already done previously."; } # this case happens only with the -i switch
			default { warn "Unknown error (malformed AVC stream?), aborting level change."; }
		}
		close H264FILE;
	}
	
	## chapters ###
	if (@$chapters) {
		print "Writing chapter file…\t";
		my $chapfile = $movie{basename} . ".chapters";
		open CHAPTERS, ">:encoding(UTF-8)", $chapfile;
		for (@$chapters) {
			print CHAPTERS "$$_{start} $$_{name}\n";
		}
		close CHAPTERS;
		push @unlink, $chapfile unless $opts{i};
		say "done.";
		
		$mp4boxcmd .= "-chap \"$chapfile\" ";
	}
	
	my $mp4file = $opts{o} ? $opts{o} : "$movie{basename}.mp4";
	$mp4boxcmd .= "\"$mp4file\"";
	
	
	$opts{e} && editVar(\$mp4boxcmd);
	$opts{v} && say "Now about to execute \n\t$mp4boxcmd";
	system $mp4boxcmd;
}
say "\n= All done! =";

say "Deleting temporary files...";
unlink foreach @unlink;


sub editVar {
my $ref = shift;

my $TEMPFILE = new File::Temp();
my $tempname = $TEMPFILE->filename;
binmode($TEMPFILE, ":encoding(UTF-8)");
print $TEMPFILE $$ref;
print $TEMPFILE <<EOF;

#######################################################################
# You can now edit above line. Only the first line will be processed. #
#######################################################################
EOF

# since we requested also a filename, it won’t get deleted right away FIXME kludgy
$TEMPFILE->close();

system("$EDITOR", $tempname);
open TF, '<:encoding(UTF-8)', $tempname;
$$ref = <TF>; # first line only
close TF;
}
